package org.farm2.service.lucene.doc;

import com.farm2.lucene.document.FarmDocument;
import org.farm2.service.lucene.doc.enums.DocumentTitleEnum;
import org.farm2.service.lucene.doc.enums.IndexModelEnum;

public class EmbeddingIndexDocument {
    private FarmDocument doc = new FarmDocument();

    public EmbeddingIndexDocument(String fileId, String knowId, String text, String title, int index, int size, float[] vector) {
        doc.addString(DocumentTitleEnum.id.name(), fileId + "-" + index);
        doc.addString(DocumentTitleEnum.model.name(), IndexModelEnum.EMBEDDING.name());
        doc.addString(DocumentTitleEnum.knowId.name(), knowId);
        doc.addText(DocumentTitleEnum.title.name(), title);
        doc.addString(DocumentTitleEnum.fileId.name(), fileId);
        doc.addText(DocumentTitleEnum.text.name(), text);
        doc.addText(DocumentTitleEnum.embedding_index.name(), String.valueOf(index));
        doc.addText(DocumentTitleEnum.embedding_size.name(), String.valueOf(size));
        doc.addVector(DocumentTitleEnum.embedding_vector.name(), vector);
    }

    public FarmDocument getDoc() {
        return doc;
    }
}
